package avltree;


public class AVLTree {

    static class TreeNode{
        public int val;
        public int bf;// 平衡因子  右树高度 - 左树高度
        public TreeNode left;
        public TreeNode right;
        public TreeNode parent;
        public TreeNode(int val) {
            this.val = val;
        }
    }
    public TreeNode root;// 根节点

    public boolean insert(int val){
        TreeNode node = new TreeNode(val);
        if (root == null) {
            root = node;
            return true;
        }

        TreeNode parent = null;
        TreeNode cur = root;
        while (cur != null){
            if (cur.val > val) {
                // 在左子树中继续找
                parent = cur;
                cur = cur.left;
            }else if (cur.val < val){
                // 在右子树中继续找
                parent = cur;
                cur = cur.right;
            }else {
                return false;// 找到节点,插不进去了
            }
        }
        // 找到对应位置, 插入节点
        if (parent.val > val){
            parent.left = node;
        }else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;
        // 调整平衡因子
        while (parent != null){
            if (cur == parent.right){
//                右树高度增加, 平衡因子++
                parent.bf++;
            }else if (cur == parent.left){
//                左树高度增加, 平衡因子--
                parent.bf--;
            }

            // 检查平衡因子的值
            if (parent.bf == 0){
                return true;
            }else if (parent.bf == 1 || parent.bf == -1){
                cur = parent;
                parent = parent.parent;
            }else {
                // 等于 2 或者 -2
                if (parent.bf == 2){// 右树高, 需要降低右树的高度
                    if (cur.bf == 1){
                        // 左旋
                        rotateLeft(parent, cur);
                    }else {
                        // cur.bf == -1
                        rotateRightLeft(parent, cur);
                    }
                }else {
                    // -2  // 左树高, 需要降低左树的高度
                    if (cur.bf == 1){
                        rotateLeftRight(parent, cur);
                    }else {
                        // cur.bf == -1
                        // 右旋
                        rotateRight(parent, cur);
                    }
                }
                break;
            }
        }

        return true;
    }

    // 右单旋
    public void rotateRight(TreeNode parent, TreeNode cur){
        TreeNode subLeft = cur;
        TreeNode subRight = cur.right;
        TreeNode pParent = parent.parent;

        subLeft.right = parent;
        parent.left = subRight;

        parent.parent = subLeft;
        if (subRight != null){
            subRight.parent = parent;
        }

        if (pParent == null){
            // parent为根节点
            root = subLeft;
            subLeft.parent = null;
        }else {
            if (parent == pParent.left){
                pParent.left = subLeft;
            }else {
                pParent.right = subLeft;
            }
            subLeft.parent = pParent;
        }

        // 调整平衡因子
        subLeft.bf = 0;
        parent.bf = 0;
    }
    // 左单旋
    public void rotateLeft(TreeNode parent, TreeNode cur){
        TreeNode subLeft = cur.left;
        TreeNode subRight = cur;
        TreeNode pParent = parent.parent;

        parent.right = subLeft;
        cur.left = parent;
        parent.parent = subRight;
        if (subLeft != null){
            subLeft.parent = parent;
        }

        if (pParent == null){
            root = subRight;
            root.parent = null;
        }else {
            if (pParent.left == parent){
                pParent.left = subRight;
            }else {
                pParent.right = subRight;
            }
            subRight.parent = pParent;
        }
        // 修改平衡因子
        subRight.bf = 0;
        parent.bf = 0;
    }
    // 左右双旋 -- 先左再右
    public void rotateLeftRight(TreeNode parent, TreeNode cur){
        TreeNode subLeft = cur;
        TreeNode subRight = subLeft.right;
        int bf = subRight.bf;

        rotateLeft(cur, cur.right);
        rotateRight(parent, parent.left);

        if (bf == 1){
            subLeft.bf = -1;
            subRight.bf = 0;
            parent.bf = 0;
        }else if (bf == -1){
//            bf == -1
            subLeft.bf = 0;
            subRight.bf = 0;
            parent.bf = 1;
        }
    }
    // 右左双旋 -- 先右再左
    private void rotateRightLeft(TreeNode parent, TreeNode cur) {
        TreeNode subRight = cur;
        TreeNode subLeft = subRight.left;
        int bf = subLeft.bf;

        rotateRight(cur, cur.left);
        rotateLeft(parent, parent.right);

        if (bf == 1){
            subLeft.bf = 0;
            parent.bf = -1;
            subRight.bf = 0;
        }else if (bf == -1){
            subLeft.bf = 0;
            parent.bf = 0;
            subRight.bf = 1;
        }
    }

    public int height(TreeNode root){
        if (root == null) return 0;

        int leftHeight = height(root.left);
        int rightHeight = height(root.right);

        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }

    public boolean isBalance(TreeNode root){
        if (root == null) return true;

//        if (inorder(root)) return false;
        int left = height(root.left);
        int right = height(root.right);

        if (right - left != root.bf){
            System.out.println("平衡因子有问题！ " + root.val);
            return false;
        }

        return Math.abs(left - right) <= 1 &&
                isBalance(root.left) && isBalance(root.right);
    }
}
